package cz.data.common.database.dialect;

import com.alibaba.fastjson.JSONObject;
import cz.data.common.support.db.core.DbColumn;
import cz.data.common.support.db.core.DbTable;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.jdbc.core.PreparedStatementCreator;
import org.springframework.jdbc.core.RowMapper;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Postgre 数据库方言
 *
 */
public class PostgreDialect extends AbstractDbDialect {

    @Override
    public String tableNames(String dbName, String schemaName) {
        String schema = ObjectUtils.defaultIfNull(schemaName, "public");
        StringBuilder sb = new StringBuilder();
        sb.append("select tablename from pg_tables ");
        sb.append("where schemaname = '").append(schema).append("'");
        sb.append(" union all ");
        sb.append("select tablename from pg_views ");
        sb.append("where schemaname = '").append(schema).append("'");
        return sb.toString();
    }

    @Override
    public String columnNames(String dbName, String schemaName, String tableName) {
        String schema = ObjectUtils.defaultIfNull(schemaName, "public");
        StringBuilder sb = new StringBuilder();
        sb.append("select column_name from information_schema.columns ");
        sb.append("where table_schema = '").append(schema);
        sb.append("' and table_name = '").append(tableName).append("'");
        return sb.toString();
    }

    @Override
    public String columns(String dbName, String schemaName, String tableName) {
        String schema = ObjectUtils.defaultIfNull(schemaName, "public");
        return "select col.column_name AS COLNAME, col.ordinal_position AS COLPOSITION, col.column_default AS DATADEFAULT, col.is_nullable AS NULLABLE, col.udt_name AS DATATYPE, " +
                "col.character_maximum_length AS DATALENGTH, col.numeric_precision AS DATAPRECISION, col.numeric_scale AS DATASCALE, des.description AS COLCOMMENT, " +
                "case when t.colname is null then 0 else 1 end as COLKEY " +
                "from information_schema.columns col left join pg_description des on col.table_name::regclass = des.objoid and col.ordinal_position = des.objsubid " +
                "left join ( " +
                "select pg_attribute.attname as colname from pg_constraint inner join pg_class on pg_constraint.conrelid = pg_class.oid " +
                "inner join pg_attribute on pg_attribute.attrelid = pg_class.oid and pg_attribute.attnum = any(pg_constraint.conkey) " +
                "where pg_class.relname = '" + tableName + "' and pg_constraint.contype = 'p' " +
                ") t on t.colname = col.column_name " +
                "where col.table_catalog = '" + dbName + "' and col.table_schema = '" + schema + "' and col.table_name = '" + tableName + "' order by col.ordinal_position ";
    }

    @Override
    public String tables(String dbName, String schemaName) {
        String schema = ObjectUtils.defaultIfNull(schemaName, "public");
        StringBuilder sb = new StringBuilder();
        sb.append("select relname AS TABLENAME, cast(obj_description(relfilenode, 'pg_class') as varchar) AS TABLECOMMENT, '1' as TABLETYPE from pg_class ");
        sb.append(" where relname in (select tablename from pg_tables where schemaname = '").append(schema).append("' )");
        sb.append(" union all ");
        sb.append("select c.relname AS TABLENAME, d.description AS TABLECOMMENT, '2' as TABLETYPE ");
        sb.append(" from pg_class c left join pg_description d on c.oid = d.objoid ");
        sb.append(" where c.relname in (select viewname from pg_views where schemaname = '").append(schema).append("' )");
        return sb.toString();
    }

    @Override
    public String buildPaginationSql(String originalSql, long offset, long count) {
        StringBuilder sqlBuilder = new StringBuilder(originalSql);
        sqlBuilder.append(" LIMIT ").append(count).append(" offset ").append(offset);
        return sqlBuilder.toString();
    }

    @Override
    public RowMapper<DbColumn> columnMapper() {
        return (ResultSet rs, int rowNum) -> {
            DbColumn entity = new DbColumn();
            entity.setColName(rs.getString("COLNAME"));
            entity.setDataType(rs.getString("DATATYPE"));
            entity.setDataLength(rs.getString("DATALENGTH"));
            entity.setDataPrecision(rs.getString("DATAPRECISION"));
            entity.setDataScale(rs.getString("DATASCALE"));
            entity.setColKey("1".equals(rs.getString("COLKEY")));
            entity.setNullable("YES".equals(rs.getString("NULLABLE")));
            entity.setColPosition(rs.getInt("COLPOSITION"));
            entity.setDataDefault(rs.getString("DATADEFAULT"));
            entity.setColComment(rs.getString("COLCOMMENT"));
            return entity;
        };
    }

    @Override
    public RowMapper<DbTable> tableMapper() {
        return (ResultSet rs, int rowNum) -> {
            DbTable entity = new DbTable();
            entity.setTableName(rs.getString("TABLENAME"));
            entity.setTableComment(rs.getString("TABLECOMMENT"));
            String tableType = rs.getString("TABLETYPE");
            if (StringUtils.isNotBlank(tableType)) {
                entity.setTableType(tableType);
            }
            return entity;
        };
    }

    @Override
    public PreparedStatementCreator upsertSQL(List<JSONObject> jsonList, String table, List<String> primaryKeys, String[] insertFields, String[] updateFields) {
        StringBuilder sb = new StringBuilder();
        sb.append("INSERT INTO \"").append(table).append("\"");
        sb.append(Stream.of(insertFields).map(this::escape).collect(Collectors.joining(",", "(", ")")));
        sb.append(" VALUES");
        sb.append(Stream.of(insertFields).map(f -> "?").collect(Collectors.joining(",", "(", ")")));
        sb.append(" ON CONFLICT ").append(primaryKeys.stream().collect(Collectors.joining(",", "(", ")")));
        sb.append(" DO UPDATE SET ");
        sb.append(Stream.of(updateFields)
                .map(this::escape)
                .map(f -> f + " = ?")
                .collect(Collectors.joining(","))
        );
        String sql = sb.toString();
        return con -> {
            con.setTransactionIsolation(Connection.TRANSACTION_READ_COMMITTED);
            PreparedStatement ps = con.prepareStatement(sql);
            for (JSONObject json : jsonList) {
                int i = 0;
                for (String insertField : insertFields) {
                    setObjectOrNULL(ps, ++i, json.get(insertField));
                }
                for (String updateField : updateFields) {
                    setObjectOrNULL(ps, ++i, json.get(updateField));
                }
                ps.addBatch();
            }
            return ps;
        };
    }

    private String escape(String field) {
        return "\"" + field + "\"";
    }
}
